# -*- coding: utf-8 -*-
# =====================
# ChatBI Query Tool
#
# Author: Assistant
# Date:   2025/1/27
# =====================
import json
import os
import asyncio
import aiohttp
from typing import AsyncGenerator

from genie_tool.util.log_util import logger, timer
from genie_tool.model.document import Doc

from typing import AsyncGenerator, List

from genie_tool.model.protocal import StreamMode
from genie_tool.tool.search_component.query_process import query_decompose
from genie_tool.tool.search_component.answer import answer_question
from genie_tool.model.protocal import StreamMode
from genie_tool.util.file_util import truncate_files
from genie_tool.model.context import LLMModelInfoFactory

class ChatBiQuery:
    """ChatBI查询工具类"""
    
    def __init__(self):
        """初始化ChatBI查询工具"""
        self._module=os.getenv("CHATBI_QUERY_MODEL")
        self._url=os.getenv("CHATBI_QUERY_URL")
        self._api_key=os.getenv("CHATBI_QUERY_API_KEY")
        self._space_id=os.getenv("CHATBI_QUERY_SPACE_ID")
        self._user=os.getenv("CHATBI_QUERY_USER")
        self._is_stream=os.getenv("CHATBI_QUERY_IS_STREAM")
        self._deep_think=os.getenv("CHATBI_QUERY_DEEP_THINKING_MODE")
        self._show_sql=os.getenv("CHATBI_QUERY_SHOW_SQL")
        self._timeout = int(os.getenv("CHATBI_QUERY_TIMEOUT", "8000"))
        self._max_retries = int(os.getenv("CHATBI_QUERY_MAX_RETRIES", "3"))


        self.searched_queries: List[str] = []
        self.current_docs: List[Doc] = []

    def search_docs_str(self, model: str = None) -> str:
        """将当前已检索文档转换为 HTML 片段文本（供回答拼接）"""
        current_docs_str = ""
        max_tokens = LLMModelInfoFactory.get_context_length(model)
        truncate_docs = truncate_files(self.current_docs, max_tokens=int(max_tokens * 0.8)) if model else self.current_docs
        for i, doc in enumerate(truncate_docs, start=1):
            current_docs_str += f"文档编号〔{i}〕. \n{doc.to_html()}\n"
        return current_docs_str
    
    def _format_chart_content(self, dimensions, metrics, sql_data):
        """
        格式化图表内容为可读文本
        根据用户要求的格式：dimensions[0][name]：sqlData[dimensions[0][name]][0]，dimensions[1][name]：sqlData[dimensions[1][name]][0]，dimensions[2][name]：sqlData[dimensions[2][name]][0]的 metrics[0][name]为：sqlData[metrics[0][name]][0]
        """
        if not sql_data:
            return str(sql_data)
        
        try:
            # 获取维度和指标名称
            dim_names = [dim.get("name", "") for dim in dimensions if isinstance(dim, dict)] if dimensions else []
            metric_names = [metric.get("name", "") for metric in metrics if isinstance(metric, dict)] if metrics else []
            
            # 如果既没有维度也没有指标，直接返回原始数据
            if not dim_names and not metric_names:
                return str(sql_data)
            
            # 构建格式化文本
            formatted_lines = []
            
            # 获取数据长度（从维度或指标中获取）
            data_length = 0
            # 优先从维度获取长度
            for dim_name in dim_names:
                if dim_name in sql_data and isinstance(sql_data[dim_name], list):
                    data_length = len(sql_data[dim_name])
                    break
            
            # 如果维度为空，从指标获取长度
            if data_length == 0:
                for metric_name in metric_names:
                    if metric_name in sql_data and isinstance(sql_data[metric_name], list):
                        data_length = len(sql_data[metric_name])
                        break
            
            # 为每行数据生成格式化文本
            for i in range(data_length):  
                line_parts = []
                
                # 添加维度信息（如果存在）
                for dim_name in dim_names:
                    if dim_name in sql_data and isinstance(sql_data[dim_name], list) and i < len(sql_data[dim_name]):
                        value = sql_data[dim_name][i]
                        line_parts.append(f"{dim_name}：{value}")
                
                # 添加指标信息（如果存在）
                for metric_name in metric_names:
                    if metric_name in sql_data and isinstance(sql_data[metric_name], list) and i < len(sql_data[metric_name]):
                        value = sql_data[metric_name][i]
                        line_parts.append(f"{metric_name}为：{value}")
                
                if line_parts:
                    formatted_lines.append("，".join(line_parts) + "；")
                       
            return "\n".join(formatted_lines) if formatted_lines else str(sql_data)
            
        except Exception as e:
            logger.error(f"格式化图表内容失败: {e}")
            return str(sql_data)
        
    @timer()
    async def run(
        self,
        query: str,
        request_id: str = None,
        max_loop: int = 1,
        stream: bool = False,
        stream_mode: StreamMode = StreamMode(),
        *args,
        **kwargs
    ) -> AsyncGenerator[str, None]:
        """执行ChatBI查询（非流式）"""

        current_loop = 1
        while current_loop <= max_loop:
            logger.info(f"{request_id} 开始ChatBI查询: {query}")
            # 查询分解
            sub_queries = await query_decompose(query=query)
    
            yield json.dumps({
                "requestId": request_id,
                "query": query,
                "searchResult": {"query": sub_queries, "docs": [[]] * len(sub_queries)},
                "isFinal": False,
                "messageType": "extend"
            }, ensure_ascii=False)
    
            await asyncio.sleep(0.1)
    
            # 去除已经检索过的query
            sub_queries = [sub_query for sub_query in sub_queries
                           if sub_query not in self.searched_queries]
            # 并行搜索并去重
            docs_list: List[List[Doc]] = []
            for sub in sub_queries:
                docs = await self._execute_query(query=sub, request_id=request_id)
                docs_list.append(docs)

            # 推送检索阶段结果
            truncate_len = int(os.getenv("SINGLE_PAGE_MAX_SIZE", 200))
            yield json.dumps(
                {
                    "requestId": request_id,
                    "query": query,
                    "searchResult": {
                        "query": sub_queries,
                        "docs": [[d.to_dict(truncate_len=truncate_len) for d in docs_l] for docs_l in docs_list]
                    },
                    "isFinal": False,
                    "messageType": "search"
                }, ensure_ascii=False)
                
            # 更新上下文
            self.current_docs.extend([d for group in docs_list for d in group])
            self.searched_queries.extend(sub_queries)

            # 如果是最后一轮，直接跳出
            if current_loop == max_loop:
                break
        # 生成最终答案
        answer = ""
        acc_content = ""
        acc_token = 0
        async for chunk in answer_question(
                query=query, search_content=self.search_docs_str(os.getenv("SEARCH_ANSWER_MODEL"))
        ):
            if stream:
                if acc_token >= stream_mode.token:
                    yield json.dumps({
                        "requestId": request_id,
                        "query": query,
                        "searchResult": {
                            "query": [],
                            "docs": [],
                        },
                        "answer": acc_content,
                        "isFinal": False,
                        "messageType": "report"
                    }, ensure_ascii=False)
                    acc_content = ""
                    acc_token = 0
                acc_content += chunk
                acc_token += 1
            answer += chunk
        if stream and acc_content:
            yield json.dumps({
                "requestId": request_id,
                "query": query,
                "searchResult": {
                    "query": [],
                    "docs": [],
                },
                "answer": acc_content,
                "isFinal": False,
                "messageType": "report"
            }, ensure_ascii=False)
        yield json.dumps({
                "requestId": request_id,
                "query": query,
                "searchResult": {
                    "query": [],
                    "docs": [],
                },
                "answer": "" if stream else answer,
                "isFinal": True,
                "messageType": "report"
            }, ensure_ascii=False)

            
    async def _execute_query(self, query: str, request_id: str) -> List[Doc]:
        """执行查询请求"""
        assert self._url and self._api_key and self._space_id and self._user,  "ChatBI配置缺失"
        url = self._url
        api = self._api_key
        headers = {
            "Authorization": f"Bearer {api}",
            "Content-Type": "application/json"
        }
        # 构建请求数据
        request_data = {
                         "model": "chatbi",
                         "stream": self._is_stream,
                         "deepThinkingMode": self._deep_think,
                         "showSql": self._show_sql,
                         "messages": [
                           {
                               "role": "user",
                               "user": self._user,
                               "content": query,
                               "spaceId": self._space_id
                           }
                         ]
                       }
        retry_count = 0
        while retry_count < self._max_retries:
            try:
                async with aiohttp.ClientSession( timeout=aiohttp.ClientTimeout(total=self._timeout)) as session:
                    logger.info(f"开始ChatBI查询: {query}")
                    # 发送POST请求
                    logger.info(f"调用ChatBI查询 url:\n{url}")
                    logger.info(f"调用ChatBI查询 headers:\n{headers}")
                    logger.info(f"调用ChatBI查询 request_data:\n{request_data}")
                    async with session.post(
                        url,
                        json=request_data,
                        headers=headers,
                        timeout=8000
                    ) as response:
                        logger.info(f"ChatBI查询响应状态码: {response.status}")
                        logger.info(f"ChatBI查询响应体：{await response.text()}")
                        response_data = await response.json(content_type=None)
                        if response.status != 200:
                            error_msg = f"ChatBI查询失败，状态码: {response.status}"
                            logger.error(f"{request_id} {error_msg}")
                            # 返回空的Doc列表，因为这是_execute_query方法，应该返回List[Doc]
                            return []
                            
                        # 处理非流式响应
                        logger.info(f"ChatBI查询原始响应体 response_data：{response_data}")
                        
                        # 检查response_data是否为None
                        if not response_data:
                            error_msg = f"ChatBI查询响应中response_data为空: {response_data}"
                            logger.error(f"{request_id} {error_msg}")
                            return []
                        
                        all_records = response_data.get("choices", [])
                        logger.info(f"ChatBI查询响应体 choices：{all_records}")
                        docs: List[Doc] = []
                        for record in all_records:
                            message = record.get("message", {})
                            logger.info(f"ChatBI查询响应体 message：{message}")
                            content_str = message.get("content", "")
                            logger.info(f"ChatBI查询响应体 content_str：{content_str}")
                            
                            # 解析JSON字符串为字典对象
                            try:
                                content = json.loads(content_str) if content_str else {}
                            except (json.JSONDecodeError, TypeError) as e:
                                logger.error(f"解析content JSON失败: {e}")
                                content = {}
                            
                            chart_info = content.get("chartInfo", {})
                            chart_title = chart_info.get("chartTitle", "")
                            logger.info(f"ChatBI查询响应体 chart_title：{chart_title}")
                            link = content.get("chatbiDetailUrl", "")
                            logger.info(f"ChatBI查询响应体 link：{link}")
                            query_sql = content.get("querySql", "")
                            logger.info(f"ChatBI查询响应体 query_sql：{query_sql}")
                            dimensions = chart_info.get("dimensions", {})
                            logger.info(f"ChatBI查询响应体 dimensions：{dimensions}")
                            metrics = chart_info.get("metrics", {})
                            logger.info(f"ChatBI查询响应体 metrics：{metrics}")
                            # 从chartInfo中获取sqlData
                            query_data = chart_info.get("sqlData", {})
                            logger.info(f"ChatBI查询响应体 query_data：{query_data}")
                            
                            # 格式化内容为可读文本
                            formatted_content = self._format_chart_content(dimensions, metrics, query_data)
                            logger.info(f"ChatBI查询响应体 formatted_content：{formatted_content}")
                            
                            docs.append(Doc(
                                doc_type="web_page",
                                content=formatted_content,
                                title=chart_title,
                                link=link,
                                data={
                                    "query_sql": query_sql,
                                    "demensions": dimensions,
                                    "metrics": metrics
                                }
                            ))
                    return docs   
                
                retry_count += 1
                
            except asyncio.TimeoutError:
                retry_count += 1
                error_msg = f"ChatBI查询超时，重试次数: {retry_count}/{self._max_retries}"
                logger.warning(f"{request_id} {error_msg}")
                
                if retry_count >= self._max_retries:
                    # 返回空的Doc列表，因为这是_execute_query方法，应该返回List[Doc]
                    return []
                    
                await asyncio.sleep(3)  # 重试前等待3秒
                
            except Exception as e:
                retry_count += 1
                error_msg = f"ChatBI查询异常: {str(e)}, 重试次数: {retry_count}/{self._max_retries}"
                logger.error(f"{request_id} {error_msg}")
                
                if retry_count >= self._max_retries:
                    # 返回空的Doc列表，因为这是_execute_query方法，应该返回List[Doc]
                    return []
                    
                await asyncio.sleep(3)  # 重试前等待3秒
                
        # 返回空的Doc列表，因为这是_execute_query方法，应该返回List[Doc]
        return []

